{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PjRuzfwyImeC"
   },
   "source": [
    "# Getting Started with NeMo Agent Toolkit\n",
    "\n",
    "In this notebook, we walk through the basics of using NVIDIA NeMo Agent toolkit (NAT), from installation all the way to creating and running a simple workflow. The intention of this notebook is to get new NAT users up and running with a high level understanding of our YAML-first approach, while gaining some intuition towards how NAT workflows can quickly be embedded into your projects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "- [0) Setup](#setup)\n",
    "  - [0.1) Prerequisites](#prereqs)\n",
    "  - [0.2) API Keys](#api-keys)\n",
    "  - [0.3) Installing NeMo Agent Toolkit](#installing-nat)\n",
    "- [1) Creating Your First Workflow](#creating-your-first-workflow)\n",
    "  - [1.1) What is a NAT workflow?](#what-is-a-workflow)\n",
    "  - [1.2) Create your first workflow](#create-first-workflow)\n",
    "  - [1.3) Interpret your first workflow](#interpret-first-workflow)\n",
    "    - [Interpreting Directory Structure](#directory-structure)\n",
    "    - [Interpreting Configuration File](#configuration-file)\n",
    "    - [Interpreting Workflow Functions](#workflow-functions)\n",
    "    - [Tying It Together](#tying-it-together)\n",
    "- [2) Running Your First Workflow](#run-first-workflow)\n",
    "    - [2.1) Run with the CLI](#run-cli)\n",
    "    - [2.2) Run as a NAT server](#run-server)\n",
    "    - [2.3) Running NAT Embedded within Python](#run-embedded)\n",
    "- [Next Steps](#next-steps)\n",
    "\n",
    "<span style=\"color:rgb(0, 31, 153); font-style: italic;\">Note: In Google Colab use the Table of Contents tab to navigate.</span>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VFUT0d7NJrtv"
   },
   "source": [
    "<a id=\"setup\"></a>\n",
    "# 0.0) Setup\n",
    "\n",
    "<a id=\"prereqs\"></a>\n",
    "## 0.1) Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i4OTdB6wTdRZ"
   },
   "source": [
    "- **Platform:** Linux, macOS, or Windows\n",
    "- **Python:** version 3.11, 3.12, or 3.13\n",
    "- **Python Packages:** `pip`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x50XDSaAJwA4"
   },
   "source": [
    "<a id=\"api-keys\"></a>\n",
    "## 0.2) API Keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vy8oHmYkJxn6"
   },
   "source": [
    "For this notebook, you will need the following API keys to run all examples end-to-end:\n",
    "\n",
    "- **NVIDIA Build:** You can obtain an NVIDIA Build API Key by creating an [NVIDIA Build](https://build.nvidia.com) account and generating a key at https://build.nvidia.com/settings/api-keys\n",
    "\n",
    "Then you can run the cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"NVIDIA_API_KEY\" not in os.environ:\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wEOYG2b-J1ys"
   },
   "source": [
    "<a id=\"installing-nat\"></a>\n",
    "## 0.3) Installing NeMo Agent Toolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OSICVNHGGm9l"
   },
   "source": [
    "The recommended way to install NAT is through `pip` or `uv pip`.\n",
    "\n",
    "First, we will install `uv` which offers parallel downloads and faster dependency resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install uv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EBV2Gh9NIC8R"
   },
   "source": [
    "NeMo Agent toolkit can be installed through the PyPI `nvidia-nat` package.\n",
    "\n",
    "There are several optional subpackages available for NAT. The `langchain` subpackage contains useful components for integrating and running within [LangChain](https://python.langchain.com/docs/introduction/). Since LangChain will be used later in this notebook, let's install NAT with the optional `langchain` subpackage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "uv pip show -q \"nvidia-nat-langchain\"\n",
    "if [ $? -ne 0 ]; then\n",
    "    uv pip install \"nvidia-nat[langchain]\"\n",
    "else\n",
    "    echo \"nvidia-nat[langchain] is already installed\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "caScQ4VxJ8Ks"
   },
   "source": [
    "<a id=\"creating-your-first-workflow\"></a>\n",
    "# 1.0) Creating Your First Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"what-is-a-workflow\"></a>\n",
    "## 1.1) What is a NAT workflow?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l7kWJ8yeJJhQ"
   },
   "source": [
    "A [workflow](https://docs.nvidia.com/nemo/agent-toolkit/latest/workflows/about/index.html) in NeMo Agent Toolkit is a structured specification of how agents, models, tools (called functions), embedders, and other components are composed together to carry out a specific task. It defines which components are used, how they are connected, and how they behave when executing the task.\n",
    "\n",
    "NAT provides a convenient command-line interface called `nat` which is accessible in your active Python environment. It serves at the entrypoint to most toolkit functions.\n",
    "\n",
    "The `nat workflow create` command allows us to create a new workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"create-first-workflow\"></a>\n",
    "## 1.2) Create your first workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat workflow create getting_started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iSDMOrSQKtBr"
   },
   "source": [
    "<a id=\"interpret-first-workflow\"></a>\n",
    "## 1.3) Interpret your first workflow\n",
    "\n",
    "<a id=\"directory-structure\"></a>\n",
    "### Interpreting Directory Structure\n",
    "We can inspect the structure of the created **workflow directory**, which we've named `getting_started`, and contains the configuration files, source code, and data needed to define and run the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!find getting_started/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fjBICzW-K0kF"
   },
   "source": [
    "A summary of the high-level components are outlined below.\n",
    "\n",
    "* `configs` (symbolic link to `src/getting_started/configs`)\n",
    "* `data` (symbolic link to `src/getting_started/data`)\n",
    "* `pyproject.toml` Python project configuration file\n",
    "* `src`\n",
    "  * `getting_started`\n",
    "    * `__init__.py` Module init file (empty)\n",
    "    * `configs` Configuration directory for workflow specifications\n",
    "      * `config.yml` Workflow configuration file\n",
    "    * `data` Data directory for any dependent files\n",
    "    * `getting_started.py` User-defined code for workflow execution\n",
    "    * `register.py` Automatic registration of project components\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HAsjWuDSTjbC"
   },
   "source": [
    "<a id=\"configuration-file\"></a>\n",
    "### Interpreting Configuration File\n",
    "The workflow configuration file, `getting_started/configs/config.yml`, describes the operational characteristics of the entire workflow. Let's load its contents in the next cell and understand what this first workflow can do out of the box."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load getting_started/configs/config.yml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t6D026_fM-h2"
   },
   "source": [
    "The above workflow configuration has the following components:\n",
    "- a [built-in `current_datetime`](https://docs.nvidia.com/nemo/agent-toolkit/latest/api/nat/tool/datetime_tools/index.html#nat.tool.datetime_tools.current_datetime) function\n",
    "- a workflow-defined `getting_started` function\n",
    "- an LLM\n",
    "- an entrypoint workflow of a [built-in ReAct agent](https://docs.nvidia.com/nemo/agent-toolkit/latest/workflows/about/react-agent.html)\n",
    "\n",
    "By default, we create a [ReAct agent](https://docs.nvidia.com/nemo/agent-toolkit/latest/workflows/about/react-agent.html) equipped with both of the functions above. When called, the Agent decides which functions to call (if any) based on the intent of user input. The agent uses the LLM to help make reasoning decisions and then performs a subsequent action.\n",
    "\n",
    "This workflow configuration file is a YAML-serialized version of the [`Config`](https://docs.nvidia.com/nemo/agent-toolkit/latest/api/nat/data_models/config/index.html#nat.data_models.config.Config) class. Each category within the high-level configuration specifies runtime configuration settings for their corresponding components. For instance, the `workflow` category contains all configuration settings for the workflow entrypoint. This configuration file is validated as typed Pydantic models and fields. All configuration classes have validation rules, default values, and [documentation](https://docs.nvidia.com/nemo/agent-toolkit/latest/workflows/workflow-configuration.html#workflow-configuration-file) which enable type-safe configuration management, automatic schema generation, and validation across the entire plugin ecosystem.\n",
    "\n",
    "* `general` - General configuration section. Contains high-level configurations for front-end definitions.\n",
    "* `authentication` - Authentication provides an interface for defining and interacting with various authentication providers.\n",
    "* `llms` - LLMs provide an interface for interacting with LLM providers.\n",
    "* `embedders` - Embedders provide an interface for interacting with embedding model providers.\n",
    "* `retreivers` - Retrievers provide an interface for searching and retrieving documents.\n",
    "* `memory` - Configurations for Memory. Memories provide an interface for storing and retrieving.\n",
    "* `object_stores` - Object Stores provide a CRUD interface for objects and data.\n",
    "* `eval` - The evaluation section provides configuration options related to the profiling and evaluation of NAT workflows.\n",
    "* `tcc_strategies` (experimental) - Test Time Compute (TTC) strategy definitions.\n",
    "\n",
    "#### Type Safety and Validation\n",
    "\n",
    "Many components within the workflow configuration specify `_type`. This YAML key is used to indicate the type of the component so NAT can properly validate and instantiate a component within the workflow. For example, [`NIMModelConfig`](https://docs.nvidia.com/nemo/agent-toolkit/latest/api/nat/llm/nim_llm/index.html#nat.llm.nim_llm.NIMModelConfig) is a subclass of [`LLMBaseConfig`](https://docs.nvidia.com/nemo/agent-toolkit/latest/api/nat/data_models/llm/index.html#nat.data_models.llm.LLMBaseConfig) so when we specify: `_type: nim` in the configuration the toolkit knows to validate the configuration with `NIMModelConfig`.\n",
    "\n",
    "<span style=\"color:rgb(0, 31, 153); font-style: italic;\">**Note:** Not all configuration components are required. The simplest workflow configuration needs to only define <code>workflow</code>.</span>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tPLvWYvtTpNF"
   },
   "source": [
    "<a id=\"workflow-functions\"></a>\n",
    "## 1.4) Interpreting Workflow Functions\n",
    "\n",
    "Next, let's inspect the contents of the generated workflow function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load getting_started/src/getting_started/getting_started.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3H5fib-jTvwq"
   },
   "source": [
    "### Function Configuration\n",
    "\n",
    "The `GettingStartedFunctionConfig` specifies `FunctionBaseConfig` as a base class. There is also a `name` specified. This name is used by the toolkit to create a static mapping when `_type` is specified anywhere where a `FunctionBaseConfig` is expected, such as `workflow` or under `functions`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3WncUuuuTxxa"
   },
   "source": [
    "### Function Registration\n",
    "\n",
    "NeMo Agent toolkit relies on a configuration with builder pattern to define most components. For functions, `@register_function` is a decorator that must be specified to inform the toolkit that a function should be accessible automatically by name when referenced. The decorator requires that a `config_type` is specified. This is done to ensure type safety and validation.\n",
    "\n",
    "The parameters to the decorated function are always:\n",
    "\n",
    "1. the configuration type of the function component (FunctionBaseConfig)\n",
    "2. a Builder which can be used to dynamically query and get other workflow components (Builder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KI8H8IoqT0TX"
   },
   "source": [
    "### Function Implementation\n",
    "\n",
    "The core logic of the `getting_started` function is embedded as a function within the outer function registration. This is done for a few reasons:\n",
    "\n",
    "* Enables dynamic importing of libraries and modules on an as-needed basis.\n",
    "* Enables context manager-like resources within to support automatic closing of resources.\n",
    "* Provides the most flexibility to users when defining their own functions.\n",
    "\n",
    "Near the end of the function registration implementation, we `yield` a `FunctionInfo` object. `FunctionInfo` is a wrapper around any type of function. It is also possible to specify additional information such as schema and converters if your function relies on transformations.\n",
    "\n",
    "NAT relies on `yield` rather `return` so resources can stay alive during the lifetime of the function or workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XYfRqWaQBHLT"
   },
   "source": [
    "<a id=\"tying-it-together\"></a>\n",
    "### Tying It Together\n",
    "\n",
    "Looking back at the configuration file, the `workflow`'s `_type` is `getting_started`. This means that the configuration of `workflow` will be validated based on the `GettingStartedFunctionConfig` implementation.\n",
    "\n",
    "The `register.py` file tells NAT what should automatically be imported so it is available when the toolkit is loaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load getting_started/src/getting_started/register.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YxP2QC1rT9UQ"
   },
   "source": [
    "<a id=\"run-first-workflow\"></a>\n",
    "# 2.0) Running Your First Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9D7yNW7ySCaY"
   },
   "source": [
    "<a id=\"run-cli\"></a>\n",
    "## 2.1) Run with the CLI\n",
    "\n",
    "You can run a workflow by using `nat run` CLI command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file getting_started/configs/config.yml \\\n",
    "         --input \"Can you echo back my name, Will?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "628pQAJLSJHF"
   },
   "source": [
    "<a id=\"run-server\"></a>\n",
    "## 2.2) Run as a NAT server\n",
    "\n",
    "NAT provides another mechanism for running workflows through `nat serve`. `nat serve` creates and launches a REST FastAPI web server for interfacing with the toolkit as though it was an OpenAI-compatible endpoint. To learn more about all endpoints served by `nat serve`, please consult [this documentation](https://docs.nvidia.com/nemo/agent-toolkit/latest/reference/api-server-endpoints.html).\n",
    "\n",
    "<span style=\"color: red\"><i>note: If running this notebook in a cloud provider such as Google Colab, `dask` may be installed. If it is, you will first have to uninstall it via:</i></span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "skip_e2e_test"
    ]
   },
   "outputs": [],
   "source": [
    "!uv pip uninstall dask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HAk1zzpjWaTF"
   },
   "source": [
    "To start the FastAPI web server, issue the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "nat serve --config_file getting_started/configs/config.yml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gXbyoAnJSq-v"
   },
   "source": [
    "It will take several seconds for the server to be reachable. The default port for the server is `8000` with `localhost` access.\n",
    "\n",
    "Note that `--input` was not required for `nat serve`. To issue a request to the server, you can then do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# Issue a request to the background service\n",
    "curl --request POST \\\n",
    "  --url http://localhost:8000/chat \\\n",
    "  --header 'Content-Type: application/json' \\\n",
    "  --data '{\n",
    "    \"messages\": [\n",
    "        {\n",
    "          \"role\": \"user\",\n",
    "          \"content\": \"What is the current time?\"\n",
    "        }\n",
    "      ]\n",
    "    }' | jq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Terminate the process after completion\n",
    "!pkill -9 -f \"nat serve\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0jTB70LnW2it"
   },
   "source": [
    "<a id=\"run-embedded\"></a>\n",
    "## 2.3) Running NAT Embedded within Python\n",
    "\n",
    "The final way to run a NAT workflow is by embedding it into an already existing Python application or library.\n",
    "\n",
    "Consider the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile nat_embedded.py\n",
    "import asyncio\n",
    "import sys\n",
    "\n",
    "from nat.runtime.loader import load_config\n",
    "from nat.utils import run_workflow\n",
    "\n",
    "\n",
    "async def amain():\n",
    "    config = load_config(sys.argv[1])\n",
    "    query_num = 1\n",
    "    try:\n",
    "        while True:\n",
    "            query = input()\n",
    "            result = await run_workflow(config=config, prompt=query)\n",
    "            print(f\"Query {query_num}: {query}\")\n",
    "            print(f\"Result {query_num}: {result}\")\n",
    "            query_num += 1\n",
    "    except EOFError:\n",
    "        pass\n",
    "\n",
    "\n",
    "asyncio.run(amain())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "27FCs1byYlYb"
   },
   "source": [
    "Then we can run it as a normal Python program as shown below, or better yet, integrate with your existing services."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "python nat_embedded.py getting_started/configs/config.yml <<EOF\n",
    "What are you capable of doing?\n",
    "What does the 'current_datetime' tool do?\n",
    "What does the 'getting_started' tool do?\n",
    "What is the current time?\n",
    "Can you echo back my name, Evan?\n",
    "What is the current time?\n",
    "Can you echo back my name, Will?\n",
    "EOF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MEtrDDQUjSpX"
   },
   "source": [
    "<a id=\"next-steps\"></a>\n",
    "# 3.0) Next Steps\n",
    "\n",
    "If you already have agents codified and don't need NAT to bring up your first agent, we also support bringing existing agents into the NAT framework. In the next notebook of this series, `bringing_your_own_agent.ipynb`, we will walk you through adapting existing agents into NAT."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cnew_313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
